package matrix

func findDiagonalOrder(mat [][]int) []int {
	M := len(mat)
	N := len(mat[0])
	next := 0
	ans := make([]int, M*N)
	x, y := 0,0
	x1, y1 := 0,0
	isUp := true
	for y <= y1 {
		next = CollectDiaAns(mat, x, y, x1, y1, next, isUp, ans)
		isUp = !isUp
		if x<M-1 {
			x++
		} else {
			y++
		}
		if y1 < N-1 {
			y1++
		} else {
			x1++
		}
	}
	return ans
}

func CollectDiaAns(mat [][]int, x, y, x1, y1, next int, isUp bool,ans []int) int {
	if isUp {
		for x >= x1 {
			ans[next] = mat[x][y]
			x--
			y++
			next++
		}
		return next
	} else {
		for x1 <= x {
			ans[next] = mat[x1][y1]
			x1++
			y1--
			next++
		}
		return next
	}
}
